# encoding=utf-8
import time
import pandas as pd
import numpy as np
from enum_type import *
import traceback

TRADE_STATUS = {
    0: 'Start',
    1: 'PreTrade',
    2: 'CollectBid',
    3: 'Trade',
    4: 'Halt',
    5: 'Suspend',
    6: 'Break',
    7: 'PosTrade',
    8: 'EndTrade',
    9: 'STOPT'
}


def add_market_suffix(df):
    try:
        df.at[:, 'stock_code'] = [row.stock_code + ('.XSHG' if int(row.market_no) == 1 else '.XSHE') for i, row in
                                  df.iterrows()]
    except:
        print(df)
    return df


def convert_from_bar(data):
    bar_fields = ['symbol', 'datetime', 'open', 'high', 'low', 'close', 'volume', 'balance']
    data_list = data['data']
    fields = data['fields']
    d = {}
    for data in data_list:
        kv = dict(zip(fields, data))
        symbol = kv['stock_code'] + '.' + kv['primary_key']
        d[symbol] = kv
    return d


def convert_from_get_snapshot(data):
    data_list = data['data']
    fields = data['fields']
    result_d = {}
    ret_fields = ['business_balance', 'prev_settlement', 'low_px', 'trade_status', 'preclose_px', 'up_px', 'high_px',
                  'millisecond', 'open_px', 'turnover_ratio', 'iopv', 'settlement', 'down_px', 'ask_price',
                  'business_amount', 'nTimeStamp', 'bid_price', 'last_px', 'offer_grp', 'prod_code', 'amount',
                  'bid_grp']
    for data in data_list:
        symbol = data[0][fields[0].index('prod_code')] + '.' + data[0][fields[0].index('market_code')]
        result_d[symbol] = {field: data[0][fields[0].index(field)] for field in ret_fields if field in fields[0]}
        result_d['business_balance'] = data[0][fields[0].index('total_balance')]
        result_d['business_amount'] = data[0][fields[0].index('total_amount')]
    return result_d


def convert_from_snapshot(data):
    timestamp = time.strptime(data['data'][0][0][data['fields'][0].index('data_time')], '%Y-%m-%d %H:%M:%S')
    prod_code = data['data'][0][0][data['fields'][0].index('prod_code')]
    market_code = data['data'][0][0][data['fields'][0].index('market_code')]
    symbol = '%s.%s' % (prod_code, market_code)

    snapshot = {}
    snapshot[symbol] = {
        'nTimeStamp': 0,
        'prod_code': symbol,
        'open_px': 0,
        'high_px': 0,
        'low_px': 0,
        'last_px': 0,
        'preclose_px': 0,
        'business_balance': 0,
        'business_amount': 0,
        'amount': 0,
        'settlement': 0,
        'prev_settlement': 0,
        'turnover_ratio': 0,
        'trade_status': 'Stop',
        'up_px': 0,
        'down_px': 0,
        'millisecond': 0,
        'iopv': 0,
    }

    snapshot[symbol] = {
        'nTimeStamp': timestamp.tm_year * 10000000000 + timestamp.tm_mon * 100000000 + timestamp.tm_mday * 1000000 +
                      timestamp.tm_hour * 10000 + timestamp.tm_min * 100 + timestamp.tm_sec,
        'prod_code': symbol,
        'open_px': data['data'][0][0][data['fields'][0].index('open_price')],
        'high_px': data['data'][0][0][data['fields'][0].index('high_price')],
        'low_px': data['data'][0][0][data['fields'][0].index('low_price')],
        'last_px': data['data'][0][0][data['fields'][0].index('last_price')],
        'preclose_px': data['data'][0][0][data['fields'][0].index('pre_close')],
        'business_balance': data['data'][0][0][data['fields'][0].index('total_balance')],
        'business_amount': data['data'][0][0][data['fields'][0].index('total_amount')],
        'iopv': data['data'][0][0][data['fields'][0].index('iopv')],
        'amount': data['data'][0][0][data['fields'][0].index('current_amount')],
        'settlement': data['data'][0][0][data['fields'][0].index('settlement')],
        'prev_settlement': data['data'][0][0][data['fields'][0].index('prev_settlement')],
        'turnover_ratio': data['data'][0][0][data['fields'][0].index('turnover_ratio')],
        'trade_status': TRADE_STATUS[data['data'][0][0][data['fields'][0].index('trade_status')]],
        'up_px': data['data'][0][0][data['fields'][0].index('up_price')],
        'down_px': data['data'][0][0][data['fields'][0].index('down_price')],
        'millisecond': data['data'][0][0][data['fields'][0].index('millisecond')],
    }

    snapshot[symbol]['bid_grp'] = {}
    for idx in range(len(data['data'][0][1])):
        snapshot[symbol]['bid_grp'][idx + 1] = [
            data['data'][0][1][idx][data['fields'][1].index('offer_price')],
            data['data'][0][1][idx][data['fields'][1].index('offer_amount')],
            data['data'][0][1][idx][data['fields'][1].index('offer_count')]]

    snapshot[symbol]['offer_grp'] = {}
    for idx in range(len(data['data'][0][2])):
        snapshot[symbol]['offer_grp'][idx + 1] = [
            data['data'][0][2][idx][data['fields'][2].index('bid_price')],
            data['data'][0][2][idx][data['fields'][2].index('bid_amount')],
            data['data'][0][2][idx][data['fields'][2].index('bid_count')]]

    return snapshot


def convert_entrust_response(response):
    head = response['head']
    count = head['count']
    if count <= 0:
        return None
    fields = response['fields']
    if np.array(fields).ndim > 1:
        fields = fields[0]
    data_list = response['data']
    if np.array(data_list).ndim > 2:
        data_list = data_list[0]
    ret_list = [dict(zip(fields, data)) for data in data_list]
    return ret_list[0] if ret_list else None


def convert_trade_responnse(response):
    head = response['head']
    type = head['type']
    fields = response['fields']
    data_list = response['data']
    if np.array(fields).ndim > 1:
        fields = fields[0]
    if np.array(data_list).ndim > 2:
        data_list = data_list[0]

    for record in data_list:
        d = dict(zip(fields, record))
        third_remark = d['third_remark']
        if '#' in third_remark:
            d['order_id'] = third_remark[: third_remark.index('#')]
        else:
            d['order_id'] = third_remark
        return d
    return []


def convert_get_position_response(data, columns=[], security=None):
    if len(data) == 0:
        return []
    if len(data['data']) == 0:
        return []
    part = data
    if len(part) == 0 or len(part['data']) == 0:
        return []
    datas = part['data']
    fields = part['fields']
    df = pd.DataFrame(datas, columns=fields)
    # df = add_market_suffix(df)
    df['stock_code'] = df['stock_code'] + "." + df['market_no']
    df = df[df['current_amount'] > 0]
    if security is not None:
        stock_code = security[:6]
        df = df[df['stock_code'] == stock_code]
    df = df.set_index('stock_code')
    df['enable_amount'] = df['begin_amount']
    if columns:
        return df[columns]
    return df


def convert_response_to_df(data, columns=[]):
    try:
        cols = columns if columns else data['fields']
        df = pd.DataFrame(data['data'], columns=data['fields'])
        return df[cols]
    except Exception as e:
        print(e)
        traceback.print_exc()
        return pd.DataFrame()


def covert_his_entrust_data(entrust_json):
    entrust_list = []
    if len(entrust_json['data']) == 0:
        return pd.DataFrame(entrust_list,
                            columns=['stock_code', 'market_no', 'full_code', 'entrust_price', 'entrust_amount',
                                     'entrust_date', 'entrust_time', 'entrust_no', 'entrust_direction',
                                     'futures_direction', 'close_direction', 'futu_invest_type', 'price_type',
                                     'entrust_state', 'investunit_id', 'instance_id', 'withdraw_amount', 'ext_remark'])

    for index in entrust_json['data']:
        instance_id_temp = index[entrust_json["fields"].index("instance_id")]
        stock_code = index[entrust_json["fields"].index("stock_code")]
        market_no = index[entrust_json["fields"].index("market_no")]
        index_list = [stock_code,
                      market_no,
                      stock_code + '.' + market_no,
                      index[entrust_json["fields"].index("entrust_price")],
                      index[entrust_json["fields"].index("entrust_amount")],
                      index[entrust_json["fields"].index("entrust_date")],
                      index[entrust_json["fields"].index("entrust_time")],
                      index[entrust_json["fields"].index("entrust_no")],
                      index[entrust_json["fields"].index("entrust_direction")],
                      index[entrust_json["fields"].index("futures_direction")],
                      index[entrust_json["fields"].index("close_direction")],
                      index[entrust_json["fields"].index("futu_invest_type")],
                      index[entrust_json["fields"].index("price_type")],
                      index[entrust_json["fields"].index("entrust_status")],
                      index[entrust_json["fields"].index("investunit_id")],
                      index[entrust_json["fields"].index("instance_id")],
                      index[entrust_json["fields"].index("withdraw_amount")],
                      index[entrust_json["fields"].index("ext_remark")]]
        entrust_list.append(index_list)
    return pd.DataFrame(entrust_list,
                        columns=['stock_code', 'market_no', 'full_code', 'entrust_price', 'entrust_amount',
                                 'entrust_date', 'entrust_time', 'entrust_no', 'entrust_direction', 'futures_direction',
                                 'close_direction', 'futu_invest_type', 'price_type', 'entrust_state', 'investunit_id',
                                 'instance_id', 'withdraw_amount', 'ext_remark'])


def convert_get_ext_orders_response(data, columns=[]):
    entrust_list = []
    if len(data['data']) == 0:
        return pd.DataFrame(entrust_list, columns=columns)
    df = pd.DataFrame(data['data'], columns=data['fields'])
    df['stock_code'] = df['stock_code'] + '.' + df['market_no']
    if len(df):
        df = df.set_index('entrust_no')
    # df = add_market_suffix(df)
    for i, row in df.iterrows():
        df.at[i, 'entrust_state_cn'] = entrust_status_dict.get(str(row.entrust_state))
        df.at[i, 'entrust_direction'] = entrust_direction_dict.get(int(row.entrust_direction))

    if columns:
        return df[columns]
    return df


def convert_get_ext_trades_response(data, columns=[]):
    entrust_list = []
    if len(data['data']) == 0:
        return pd.DataFrame(entrust_list, columns=columns)
    df = pd.DataFrame(data['data'], columns=data['fields'])
    df['stock_code'] = df['stock_code'] + '.' + df['market_no']
    if len(df):
        df = df.set_index('deal_no')
    # df = add_market_suffix(df)
    if columns:
        return df[columns]
    return df


def convert_trade_status(status_no):
    kv = {'0': '未报',
          '1': '未报',
          '2': '待报',
          '3': '正报',
          '4': '已报',
          '5': '废单',
          '6': '部成',
          '7': '已成',
          '8': '部撤',
          '9': '已撤',
          'a': '待撤',
          'b': '未审批',
          'c': '审批拒绝',
          'd': '未审批即撤销',
          'e': '撤单待审批',
          'A': '未撤(撤单委托)',
          'B': '待撤(撤单委托)',
          'C': '正撤(撤单委托)',
          'D': '撤认(撤单委托)',
          'E': '撤废(撤单委托)',
          'F': '已撤(撤单委托)'}
    return kv.get(status_no, '未知')


if __name__ == '__main__':
    response = {'data': [],
                'head': {'count': 1, 'uuid': '', 'requestid': '', 'functionid': 10002, 'timestamp': '20211029160222',
                         'user': '846e1fae-f6e9-47c7-acde-d6a6af757343', 'guid': '846e1fae-f6e9-47c7-acde-d6a6af757343',
                         'type': 1, 'index': 1}, 'fields': []}
    convert_entrust_response(response)
    response = {'head': {'timestamp': '20210129152425', 'type': 1, 'index': 1, 'functionid': 10002, 'count': 1,
                         'user': '7f0b9e05-ae52-4d39-953a-b24902e28f68',
                         'guid': 'f1228fae-aa2e-423d-a7fd-c3cea8203041'}, 'data': [
        [0, '', 5035, 0, 325, 428, 0, 0, 0, None, None, '0', None, 'None', None, None, 'System.Object']],
                'fields': ['error_no', 'error_info', 'entrust_no', 'batch_no', 'investunit_id', 'instance_id',
                           'capitalunit_id', 'strategy_flag_id', 'strategy_req_id', 'success_flag', 'trace_i_d',
                           'trace_id2', 'primary_key', 'update_type', 'topic_name', 'table_name', 'read_write_obj']}
    ret_list = convert_entrust_response(response)
    for ret in ret_list:
        print(ret)

    response = {'head': {'timestamp': '20210129152425', 'type': 2, 'index': 1, 'functionid': 20001, 'count': 1,
                         'user': '7f0b9e05-ae52-4d39-953a-b24902e28f68',
                         'guid': 'b1106427-92e3-43e2-a345-6011cb86bc23'}, 'data': [[[5035, 100.0, 0.0, '1', '300345',
                                                                                     'XSHE', 325, 428, '', None, 'BUY',
                                                                                     None,
                                                                                     'cf49e3b6-613d-11eb-b455-dc1ba1d7c255#自动交易',
                                                                                     'A', 0.0]]], 'fields': [
        ['entrust_no', 'amount', 'price', 'status', 'stock_code', 'market_no', 'investunit_id', 'instance_id',
         'invest_type', 'position_type', 'entrust_direction', 'futures_direction', 'third_remark', 'style',
         'filled_amount']]}
    for ret in convert_trade_responnse(response):
        print(ret)
